// "Copyright [2021] <Copyright Shengkai Lin>"

#include <memory>
#include "route.h"

using namespace omnetpp;

namespace ai_sim {

void RoceClusterRouteAlgorithm::generateRoutes(
    taskboard& tmp_taskboard, TopoManager * topos) {
for (auto& routing_stage : tmp_taskboard.stage_info) {
// Loop in different stages

for (auto& comm_pair : routing_stage.comm_pairs) {
    topos->getDataplaneTopo()->calculateWBSSPathsTo(
        comm_pair.second);
    auto tmp_node = comm_pair.first;

    int first_id = topos->getWorkersTopo()->getNodeId(tmp_node);
    int second_id = topos->getWorkersTopo()->getNodeId(comm_pair.second);
    int gate_id_A = -1, gate_id_B = -1, gate_id_C = -1, gate_id_D = -1;

    vector<topoLink*> stage_tmp_links;
    EV << "--------------" << "\n";
    EV << "Src is " <<
        tmp_node->getModule()->getFullPath() << "\n";
    EV << "Dest is " << comm_pair.second->getModule()->
        getFullPath() << std::endl;
    EV << "Dist is " << tmp_node->getDistanceToTarget() << std::endl;
    if (tmp_node->getNumPaths() == 0) {
        throw omnetpp::cRuntimeError("CAN NOT ALLOC ANY ROUTES, CHECK THE TASK");
        return;
    }

    // First, set the output port of the first gpu.
    for (auto& flow : routing_stage.flows_on_gpu[first_id]) {
        if (flow.destIdx == second_id) {
            flow.outportId = tmp_node->getoutPathtoDest(0)->getLocalGateId();
        }
    }
    // Second, Start alloc routes for switches

    // Disable the specified links
    tmp_node->getoutPathtoDest(0)->disable();
    tmp_node->getinPathfromDest(0)->disable();
    // Record the specfied links
    routing_stage.used_links.push_back(tmp_node->getoutPathtoDest(0));
    routing_stage.used_links.push_back(tmp_node->getinPathfromDest(0));

    stage_tmp_links.push_back(tmp_node->getoutPathtoDest(0));
    stage_tmp_links.push_back(tmp_node->getinPathfromDest(0));

    gate_id_A = tmp_node->getoutPathtoDest(0)->getRemoteGateId();
    gate_id_D = tmp_node->getinPathfromDest(0)->getRemoteGateId();

    tmp_node = tmp_node->getoutPathtoDest(0)->getRemoteNode();

    while (tmp_node->getDistanceToTarget() > 0) {
        gate_id_B = tmp_node->getoutPathtoDest(0)->getLocalGateId();
        gate_id_C = tmp_node->getinPathfromDest(0)->getLocalGateId();
        RouteInfo tmp_route_f = {
            .inPort = gate_id_A,
            .outPort = gate_id_B,
            .srcIdx = first_id,
            .destIdx = second_id
        };
        RouteInfo tmp_route_b = {
            .inPort = gate_id_C,
            .outPort = gate_id_D,
            .srcIdx = second_id,
            .destIdx = first_id
        };
        routing_stage.routes_on_switch[tmp_node->getModuleId()]
            .push_back(tmp_route_f);
        // Push back depended routes for flows
        for (auto& flow : routing_stage.flows_on_gpu[first_id]) {
            if (flow.destIdx == second_id) {
                flow.dependRoute[tmp_node->getModuleId()].
                    push_back(tmp_route_f);
            }
        }
        routing_stage.routes_on_switch[tmp_node->getModuleId()]
            .push_back(tmp_route_b);
        // Push back depended routes for flows
        for (auto& flow : routing_stage.flows_on_gpu[second_id]) {
            if (flow.destIdx == first_id) {
                flow.dependRoute[tmp_node->getModuleId()].
                    push_back(tmp_route_b);
            }
        }

        EV << "We are in " <<
            tmp_node->getModule()->getFullPath() << std::endl;
        EV << "Dest is " << comm_pair.second->getModule()->
            getFullPath() << std::endl;
        EV << tmp_node->getDistanceToTarget() << " hops to go\n";

        routing_stage.used_links.push_back(tmp_node->getoutPathtoDest(0));
        routing_stage.used_links.push_back(tmp_node->getinPathfromDest(0));

        stage_tmp_links.push_back(tmp_node->getoutPathtoDest(0));
        stage_tmp_links.push_back(tmp_node->getinPathfromDest(0));

        gate_id_A = tmp_node->getoutPathtoDest(0)->getRemoteGateId();
        gate_id_D = tmp_node->getinPathfromDest(0)->getRemoteGateId();

        tmp_node = tmp_node->getoutPathtoDest(0)->getRemoteNode();
    }

    for (topoLink* link : stage_tmp_links) {
            link->enable();
        }
        
    // Finally, update the output_port for the backword flows
    for (auto& flow : routing_stage.flows_on_gpu[second_id]) {
        if (flow.destIdx == first_id) {
            flow.outportId = gate_id_D;
        }
    }
}
}
}
}  // namespace ai_sim
